import tushare as ts
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import os
import statsmodels.api as sm
from scipy import stats
from datetime import datetime, timedelta
from analysis_engine import SectorAnalyzer

class FactorValidator:
    def __init__(self):
        ts.set_token("8f0c9f36fea80fe9cc6a0d480a87ba58c48a91182d0f7562e9ac9706")
        self.pro = ts.pro_api()
        self.analyzer = SectorAnalyzer()
        ts.set_token("8f0c9f36fea80fe9cc6a0d480a87ba58c48a91182d0f7562e9ac9706")
        self.pro = ts.pro_api()
        
    def get_future_returns(self, trade_date, n_days=5):
        """
        计算从trade_date开始未来n_days的行业收益率
        """
        # 获取未来交易日
        end_date = datetime.strptime(trade_date, '%Y%m%d')
        future_date = end_date + timedelta(days=n_days*2)  # 留出足够的日历日以获取足够的交易日
        trade_dates = self.analyzer.get_trade_dates(trade_date, future_date.strftime('%Y%m%d'))
        
        if len(trade_dates) <= n_days:
            print(f"无法获取足够的未来交易日数据")
            return None
            
        future_date = trade_dates[n_days]  # 获取未来第n个交易日
        
        # 获取当前日期和未来日期的行业指数数据
        # 这里假设使用申万一级行业指数，实际使用时可能需要调整
        try:
            # 获取当前日期的行业数据
            current_data = self.analyzer.get_stock_data(trade_date)
            future_data = self.analyzer.get_stock_data(future_date)
            
            # 计算行业收益率
            # 这里简化处理，实际应该使用行业指数或成分股加权收益
            # 获取当前日期和未来日期的股票价格
            current_prices = {}
            future_prices = {}
            
            # 获取当前日期的股票价格
            stocks = list(set(current_data['ts_code']))
            if len(stocks) > 0:
                current_price_data = self.pro.daily(trade_date=trade_date, fields='ts_code,close')
                current_prices = dict(zip(current_price_data['ts_code'], current_price_data['close']))
                
                future_price_data = self.pro.daily(trade_date=future_date, fields='ts_code,close')
                future_prices = dict(zip(future_price_data['ts_code'], future_price_data['close']))
            
            # 计算每个行业的收益率
            sector_returns = {}
            for sector in current_data['sw_industry'].unique():
                sector_stocks = current_data[current_data['sw_industry'] == sector]['ts_code']
                sector_returns_list = []
                
                for stock in sector_stocks:
                    if stock in current_prices and stock in future_prices:
                        stock_return = (future_prices[stock] / current_prices[stock]) - 1
                        sector_returns_list.append(stock_return)
                
                if sector_returns_list:
                    # 使用等权平均计算行业收益
                    sector_returns[sector] = np.mean(sector_returns_list)
                else:
                    sector_returns[sector] = np.nan
            
            return pd.Series(sector_returns)
            
        except Exception as e:
            print(f"计算未来收益率失败: {str(e)}")
            return None
    
    def calculate_ic(self, factor_scores, future_returns):
        """
        计算因子IC值（Information Coefficient）
        增强版IC计算，使用Spearman秩相关系数并添加稳健性检查
        """
        # 合并因子得分和未来收益
        if future_returns is None or len(future_returns) == 0:
            return None
            
        # 确保输入为Series类型
        if not isinstance(factor_scores, pd.Series):
            factor_scores = pd.Series(factor_scores)
        if not isinstance(future_returns, pd.Series):
            future_returns = pd.Series(future_returns)
            
        # 对齐索引并合并数据
        merged_data = pd.DataFrame({
            'factor': factor_scores,
            'returns': future_returns
        }).dropna()
        
        # 增强数据量检查
        if len(merged_data) < 5:  # 至少需要5个有效数据点
            return None
            
        # 检查因子得分和收益的方差
        if merged_data['factor'].std() < 1e-6 or merged_data['returns'].std() < 1e-6:
            return None
            
        try:
            # 计算Spearman秩相关系数，添加异常处理
            ic = merged_data['factor'].corr(merged_data['returns'], method='spearman')
            
            # 添加IC值合理性检查
            if np.isnan(ic) or abs(ic) > 1.0:
                return None
                
            return ic
        except Exception as e:
            print(f"计算IC值时发生错误: {str(e)}")
            return None
    
    def run_t_test(self, ic_series):
        """
        对IC序列进行t检验，判断是否显著大于0
        """
        if ic_series is None or len(ic_series) < 10:  # 至少需要10个IC值
            return None, None
            
        # 去除缺失值
        ic_series = ic_series.dropna()
        
        # 计算均值和标准差
        ic_mean = ic_series.mean()
        ic_std = ic_series.std()
        
        # 计算t值和p值
        t_stat = ic_mean / (ic_std / np.sqrt(len(ic_series)))
        p_value = 1 - stats.t.cdf(t_stat, len(ic_series)-1)
        
        return t_stat, p_value
    
    def run_quantile_backtest(self, factor_name, start_date, end_date, n_groups=5, n_days=5):
        """
        按因子得分将行业分组，回测各组合的表现
        """
        # 获取交易日序列
        trade_dates = self.analyzer.get_trade_dates(start_date, end_date)
        
        # 存储各组合的收益率序列
        group_returns = {i: [] for i in range(1, n_groups+1)}
        dates = []
        
        for date in trade_dates:
            try:
                # 获取当日因子得分
                sector_stats, quantiles = self.analyzer.calculate_indicators(date)
                current_data = self.analyzer.get_stock_data(date)
                final_scores = self.analyzer.calculate_scores(sector_stats, quantiles, current_data)
                
                # 获取特定因子得分
                if factor_name == 'density':
                    factor_scores = final_scores.set_index('sw_industry')['density_score']
                elif factor_name == 'consecutive':
                    factor_scores = final_scores.set_index('sw_industry')['consecutive_score']
                elif factor_name == 'historical':
                    factor_scores = final_scores.set_index('sw_industry')['historical_score']
                elif factor_name == 'strength':
                    factor_scores = final_scores.set_index('sw_industry')['strength_score']
                else:
                    factor_scores = final_scores.set_index('sw_industry')['score']
                
                # 获取未来收益率
                future_returns = self.get_future_returns(date, n_days)
                
                if future_returns is not None and len(factor_scores) > 0:
                    # 合并因子得分和未来收益
                    merged_data = pd.DataFrame({
                        'factor': factor_scores,
                        'returns': future_returns
                    }).dropna()
                    
                    if len(merged_data) >= n_groups:  # 确保有足够的数据进行分组
                        # 按因子得分分组
                        merged_data['group'] = pd.qcut(merged_data['factor'], n_groups, labels=range(1, n_groups+1))
                        
                        # 计算各组合的平均收益率
                        group_avg_returns = merged_data.groupby('group')['returns'].mean()
                        
                        # 存储各组合的收益率
                        for group, ret in group_avg_returns.items():
                            group_returns[group].append(ret)
                        
                        dates.append(date)
            except Exception as e:
                print(f"日期{date}处理失败: {str(e)}")
                continue
        
        # 计算累积收益率
        cumulative_returns = {}
        for group, returns in group_returns.items():
            if returns:  # 确保有收益率数据
                cumulative_returns[group] = (1 + pd.Series(returns)).cumprod() - 1
            else:
                cumulative_returns[group] = pd.Series([])
        
        return cumulative_returns, dates
    
    def run_parameter_sensitivity(self, factor_name, start_date, end_date, param_values, n_days=5):
        """
        参数敏感性测试
        """
        # 存储不同参数下的IC值
        param_ic_values = {param: [] for param in param_values}
        
        # 获取交易日序列
        trade_dates = self.analyzer.get_trade_dates(start_date, end_date)
        
        for date in trade_dates[:30]:  # 为了效率，只使用部分日期
            try:
                # 获取未来收益率
                future_returns = self.get_future_returns(date, n_days)
                
                if future_returns is not None:
                    for param in param_values:
                        # 这里根据不同的因子和参数值计算因子得分
                        # 实际实现时需要根据具体因子调整计算逻辑
                        if factor_name == 'density':
                            # 例如，调整涨停密度的计算窗口
                            sector_stats, _ = self.analyzer.calculate_indicators(date)
                            factor_scores = sector_stats.set_index('sw_industry')['limit_density'] * param
                        elif factor_name == 'historical':
                            # 调整历史分位值的计算窗口
                            # 这里简化处理，实际应该重新计算分位值
                            sector_stats, quantiles = self.analyzer.calculate_indicators(date)
                            current_data = self.analyzer.get_stock_data(date)
                            final_scores = self.analyzer.calculate_scores(sector_stats, quantiles, current_data)
                            factor_scores = final_scores.set_index('sw_industry')['historical_score'] * param / 100
                        else:
                            # 其他因子的参数敏感性测试
                            sector_stats, quantiles = self.analyzer.calculate_indicators(date)
                            current_data = self.analyzer.get_stock_data(date)
                            final_scores = self.analyzer.calculate_scores(sector_stats, quantiles, current_data)
                            factor_scores = final_scores.set_index('sw_industry')[factor_name + '_score']
                        
                        # 计算IC值
                        ic = self.calculate_ic(factor_scores, future_returns)
                        if ic is not None:
                            param_ic_values[param].append(ic)
            except Exception as e:
                print(f"日期{date}参数敏感性测试失败: {str(e)}")
                continue
        
        # 计算各参数下的平均IC值
        avg_ic_values = {param: np.mean(ics) if ics else np.nan for param, ics in param_ic_values.items()}
        
        return avg_ic_values
    
    def run_market_state_test(self, factor_name, start_date, end_date, n_days=5):
        """
        在不同市场状态下测试因子表现
        """
        # 获取交易日序列
        trade_dates = self.analyzer.get_trade_dates(start_date, end_date)
        
        # 获取市场指数数据（以上证指数为例）
        index_data = []
        for date in trade_dates:
            try:
                # 获取上证指数数据
                df = self.pro.index_daily(ts_code='000001.SH', trade_date=date, fields='trade_date,close')
                if not df.empty:
                    index_data.append(df.iloc[0])
            except Exception as e:
                print(f"获取指数数据失败: {str(e)}")
                continue
        
        if not index_data:
            return None
            
        index_df = pd.DataFrame(index_data)
        index_df['close'] = index_df['close'].astype(float)
        
        # 计算20日收益率，用于判断市场状态
        index_df['return_20d'] = index_df['close'].pct_change(20)
        
        # 定义市场状态：牛市、熊市、震荡市
        # 简化处理：20日收益率>5%为牛市，<-5%为熊市，其他为震荡市
        index_df['market_state'] = 'flat'
        index_df.loc[index_df['return_20d'] > 0.05, 'market_state'] = 'bull'
        index_df.loc[index_df['return_20d'] < -0.05, 'market_state'] = 'bear'
        
        # 存储不同市场状态下的IC值
        market_state_ic = {'bull': [], 'bear': [], 'flat': []}
        
        for i, date in enumerate(trade_dates):
            if i >= len(index_df):
                break
                
            try:
                market_state = index_df.iloc[i]['market_state']
                
                # 获取当日因子得分
                sector_stats, quantiles = self.analyzer.calculate_indicators(date)
                current_data = self.analyzer.get_stock_data(date)
                final_scores = self.analyzer.calculate_scores(sector_stats, quantiles, current_data)
                
                # 获取特定因子得分
                if factor_name == 'density':
                    factor_scores = final_scores.set_index('sw_industry')['density_score']
                elif factor_name == 'consecutive':
                    factor_scores = final_scores.set_index('sw_industry')['consecutive_score']
                elif factor_name == 'historical':
                    factor_scores = final_scores.set_index('sw_industry')['historical_score']
                elif factor_name == 'strength':
                    factor_scores = final_scores.set_index('sw_industry')['strength_score']
                else:
                    factor_scores = final_scores.set_index('sw_industry')['score']
                
                # 获取未来收益率
                future_returns = self.get_future_returns(date, n_days)
                
                # 计算IC值
                ic = self.calculate_ic(factor_scores, future_returns)
                if ic is not None:
                    market_state_ic[market_state].append(ic)
            except Exception as e:
                print(f"市场状态测试失败: {str(e)}")
                continue
        
        # 计算各市场状态下的平均IC值
        avg_market_state_ic = {state: np.mean(ics) if ics else np.nan for state, ics in market_state_ic.items()}
        
        return avg_market_state_ic
    
    def run_risk_adjusted_returns(self, factor_name, start_date, end_date, n_days=5):
        """
        计算风险调整后收益
        """
        # 获取交易日序列
        trade_dates = self.analyzer.get_trade_dates(start_date, end_date)
        
        # 存储高分组和低分组的收益率序列
        high_group_returns = []
        low_group_returns = []
        market_returns = []
        dates = []
        
        for date in trade_dates:
            try:
                # 获取当日因子得分
                sector_stats, quantiles = self.analyzer.calculate_indicators(date)
                current_data = self.analyzer.get_stock_data(date)
                final_scores = self.analyzer.calculate_scores(sector_stats, quantiles, current_data)
                
                # 获取特定因子得分
                if factor_name == 'density':
                    factor_scores = final_scores.set_index('sw_industry')['density_score']
                elif factor_name == 'consecutive':
                    factor_scores = final_scores.set_index('sw_industry')['consecutive_score']
                elif factor_name == 'historical':
                    factor_scores = final_scores.set_index('sw_industry')['historical_score']
                elif factor_name == 'strength':
                    factor_scores = final_scores.set_index('sw_industry')['strength_score']
                else:
                    factor_scores = final_scores.set_index('sw_industry')['score']
                
                # 获取未来收益率
                future_returns = self.get_future_returns(date, n_days)
                
                # 获取市场收益率（以上证指数为例）
                try:
                    current_index = self.pro.index_daily(ts_code='000001.SH', trade_date=date, fields='close').iloc[0]['close']
                    future_date = trade_dates[trade_dates.index(date) + n_days] if trade_dates.index(date) + n_days < len(trade_dates) else None
                    if future_date:
                        future_index = self.pro.index_daily(ts_code='000001.SH', trade_date=future_date, fields='close').iloc[0]['close']
                        market_return = (future_index / current_index) - 1
                        market_returns.append(market_return)
                    else:
                        continue
                except Exception as e:
                    print(f"获取市场收益率失败: {str(e)}")
                    continue
                
                if future_returns is not None and len(factor_scores) > 0:
                    # 合并因子得分和未来收益
                    merged_data = pd.DataFrame({
                        'factor': factor_scores,
                        'returns': future_returns
                    }).dropna()
                    
                    if len(merged_data) >= 5:  # 确保有足够的数据进行分组
                        # 按因子得分分为高分组和低分组
                        high_group = merged_data.nlargest(len(merged_data) // 5, 'factor')
                        low_group = merged_data.nsmallest(len(merged_data) // 5, 'factor')
                        
                        # 计算高分组和低分组的平均收益率
                        high_group_return = high_group['returns'].mean()
                        low_group_return = low_group['returns'].mean()
                        
                        high_group_returns.append(high_group_return)
                        low_group_returns.append(low_group_return)
                        dates.append(date)
            except Exception as e:
                print(f"风险调整收益计算失败: {str(e)}")
                continue
        
        if not high_group_returns or not low_group_returns or not market_returns:
            return None
            
        # 计算高分组和低分组的累积收益率
        high_cumulative = (1 + pd.Series(high_group_returns)).cumprod() - 1
        low_cumulative = (1 + pd.Series(low_group_returns)).cumprod() - 1
        market_cumulative = (1 + pd.Series(market_returns)).cumprod() - 1
        
        # 计算高分组的最大回撤
        high_max_drawdown = self.calculate_max_drawdown(high_cumulative)
        
        # 计算高分组的夏普比率（简化版，假设无风险利率为0）
        high_returns = pd.Series(high_group_returns)
        high_sharpe = high_returns.mean() / high_returns.std() * np.sqrt(252/n_days) if high_returns.std() > 0 else np.nan
        
        # 计算高分组相对于市场的alpha和beta
        market_returns_series = pd.Series(market_returns)
        if len(high_returns) == len(market_returns_series) and len(high_returns) > 2:
            model = sm.OLS(high_returns, sm.add_constant(market_returns_series))
            results = model.fit()
            alpha = results.params[0] * 252  # 年化alpha
            beta = results.params[1]
        else:
            alpha = np.nan
            beta = np.nan
        
        return {
            'high_cumulative': high_cumulative,
            'low_cumulative': low_cumulative,
            'market_cumulative': market_cumulative,
            'max_drawdown': high_max_drawdown,
            'sharpe_ratio': high_sharpe,
            'alpha': alpha,
            'beta': beta,
            'dates': dates
        }
    
    def calculate_max_drawdown(self, returns):
        """
        计算最大回撤
        """
        cumulative = returns
        if len(cumulative) == 0:
            return np.nan
            
        # 计算当前高点
        running_max = cumulative.cummax()
        # 计算从当前高点回撤
        drawdown = (cumulative / running_max) - 1
        # 最大回撤
        max_drawdown = drawdown.min()
        
        return max_drawdown
    
    def validate_factor(self, factor_name, start_date, end_date, n_days=5):
        """
        综合验证因子有效性
        """
        print(f"开始验证因子: {factor_name}")
        
        # 1. 统计显著性检验
        print("\n1. 统计显著性检验")
        
        # 计算IC序列
        trade_dates = self.analyzer.get_trade_dates(start_date, end_date)
        ic_values = []
        valid_dates = []
        
        for date in trade_dates:
            try:
                # 获取当日因子得分
                sector_stats, quantiles = self.analyzer.calculate_indicators(date)
                current_data = self.analyzer.get_stock_data(date)
                final_scores = self.analyzer.calculate_scores(sector_stats, quantiles, current_data)
                
                # 获取特定因子得分
                if factor_name == 'density':
                    factor_scores = final_scores.set_index('sw_industry')['density_score']
                elif factor_name == 'consecutive':
                    factor_scores = final_scores.set_index('sw_industry')['consecutive_score']
                elif factor_name == 'historical':
                    factor_scores = final_scores.set_index('sw_industry')['historical_score']
                elif factor_name == 'strength':
                    factor_scores = final_scores.set_index('sw_industry')['strength_score']
                else:
                    factor_scores = final_scores.set_index('sw_industry')['score']
                
                # 获取未来收益率
                future_returns = self.get_future_returns(date, n_days)
                
                # 计算IC值
                ic = self.calculate_ic(factor_scores, future_returns)
                if ic is not None:
                    ic_values.append(ic)
                    valid_dates.append(date)
            except Exception as e:
                print(f"日期{date}处理失败: {str(e)}")
                continue
        
        if ic_values:
            ic_series = pd.Series(ic_values, index=valid_dates)
            avg_ic = ic_series.mean()
            ic_std = ic_series.std()
            ir = avg_ic / ic_std if ic_std > 0 else np.nan
            t_stat, p_value = self.run_t_test(ic_series)
            
            print(f"平均IC: {avg_ic:.4f}")
            print(f"IC标准差: {ic_std:.4f}")
            print(f"信息比率(IR): {ir:.4f}")
            if t_stat is not None and p_value is not None:
                print(f"t统计量: {t_stat:.4f}")
                print(f"p值: {p_value:.4f}")
            print(f"IC>0的比例: {(ic_series > 0).mean():.4f}")
            
            # 绘制IC序列图
            plt.figure(figsize=(10, 6))
            plt.plot(ic_series)
            plt.axhline(y=0, color='r', linestyle='-')
            plt.axhline(y=avg_ic, color='g', linestyle='--')
            plt.title(f'{factor_name} IC序列 (平均IC={avg_ic:.4f})')
            plt.xlabel('日期')
            plt.ylabel('IC值')
            plt.grid(True)
            plt.tight_layout()
            
            # 保存图表
            output_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'output')
            if not os.path.exists(output_dir):
                os.makedirs(output_dir)
            plt.savefig(os.path.join(output_dir, f'{factor_name}_ic_series.png'))
            plt.close()
        else:
            print("无法计算有效的IC序列")
        
        # 2. 分位数回测
        print("\n2. 分位数回测")
        cumulative_returns, dates = self.run_quantile_backtest(factor_name, start_date, end_date, n_groups=5, n_days=n_days)
        
        if cumulative_returns and any(len(ret) > 0 for ret in cumulative_returns.values()):
            # 绘制累积收益率图
            plt.figure(figsize=(10, 6))
            for group, returns in cumulative_returns.items():
                if len(returns) > 0:
                    plt.plot(returns, label=f'Group {group}')
            plt.title(f'{factor_name} 分位数回测 (未来{n_days}日收益)')
            plt.xlabel('交易日序号')
            plt.ylabel('累积收益率')
            plt.legend()
            plt.grid(True)
            plt.tight_layout()
            
            # 保存图表
            plt.savefig(os.path.join(output_dir, f'{factor_name}_quantile_returns.png'))
            plt.close()
            
            # 计算高分组和低分组的收益差异
            if 5 in cumulative_returns and 1 in cumulative_returns and len(cumulative_returns[5]) > 0 and len(cumulative_returns[1]) > 0:
                high_low_spread = cumulative_returns[5].iloc[-1] - cumulative_returns[1].iloc[-1]
                print(f"高分组与低分组收益差异: {high_low_spread:.4f}")
        else:
            print("无法进行有效的分位数回测")
        
        # 3. 参数敏感性测试
        print("\n3. 参数敏感性测试")
        if factor_name == 'density':
            param_values = [0.8, 0.9, 1.0, 1.1, 1.2]  # 调整涨停密度的权重
        elif factor_name == 'historical':
            param_values = [70, 80, 90, 95]  # 调整历史分位值的阈值
        else:
            param_values = [0.8, 0.9, 1.0, 1.1, 1.2]  # 通用参数调整
            
        avg_ic_values = self.run_parameter_sensitivity(factor_name, start_date, end_date, param_values, n_days)
        
        if avg_ic_values and not all(np.isnan(v) for v in avg_ic_values.values()):
            # 绘制参数敏感性图
            plt.figure(figsize=(10, 6))
            params = list(avg_ic_values.keys())
            ics = [avg_ic_values[p] for p in params]
            plt.bar(range(len(params)), ics)
            plt.xticks(range(len(params)), params)
            plt.title(f'{factor_name} 参数敏感性测试 (平均IC)')
            plt.xlabel('参数值')
            plt.ylabel('平均IC')
            plt.grid(True)
            plt.tight_layout()
            
            # 保存图表
            plt.savefig(os.path.join(output_dir, f'{factor_name}_parameter_sensitivity.png'))